import contextlib
from base64 import b64decode
from io import BytesIO, StringIO
from typing import Any
from urllib.parse import ParseResult, urlparse

import paramiko

from worker.log import logger

from .base_publisher import BasePublisher


class SFTPPublisher(BasePublisher):
    REQUIRED_PARAMETERS = ("SFTP_URL",)

    def __init__(self):
        super().__init__()
        self.ftp_url = None
        self.ssh = paramiko.SSHClient()

        self.type = "SFTP_PUBLISHER"
        self.name = "SFTP Publisher"
        self.description = "Publisher for publishing to a SFTP server"

    def publish(self, publisher: dict[str, Any], product, rendered_product):
        parameters = self._extract_parameters(publisher)
        ftp_url = parameters.get("SFTP_URL")
        private_key = parameters.get("PRIVATE_KEY")

        self.set_file_name(product)
        server_config: ParseResult = urlparse(ftp_url)  # type: ignore

        if rendered_product.mime_type in ["text/plain", "text/html"]:
            data_to_upload = BytesIO(rendered_product.data)
        else:
            data_to_upload = BytesIO(b64decode(rendered_product.data))

        if not server_config:
            raise ValueError("Invalid SFTP URL")
        self.input_validation(server_config)
        if private_key:
            private_key = self.parse_private_key(private_key)

        self.upload_to_sftp(server_config, data_to_upload, private_key=private_key)

        return "SFTP Publisher Task Successful"

    def input_validation(self, server_config: ParseResult):
        if not server_config.username:
            logger.error(f"{server_config.username=}")
            raise ValueError("Username is required for SFTP")

        if server_config.scheme != "sftp":
            raise ValueError(f"Schema '{server_config.scheme}' not supported, choose 'sftp'")

    def parse_private_key(self, private_key: str) -> paramiko.PKey:
        from paramiko.ecdsakey import ECDSAKey
        from paramiko.ed25519key import Ed25519Key
        from paramiko.rsakey import RSAKey

        for pkey_class in (RSAKey, ECDSAKey, Ed25519Key):
            with contextlib.suppress(paramiko.SSHException):
                return pkey_class.from_private_key(StringIO(private_key))
        raise ValueError("Invalid private key format for SFTP")

    def upload_to_sftp(self, server_config: ParseResult, data_to_upload: BytesIO, private_key: paramiko.PKey | None = None):
        ssh_port = server_config.port or 22
        remote_path = server_config.path + self.file_name
        connect_password = None if private_key else server_config.password
        hostname = server_config.hostname
        if not hostname:
            raise ValueError("Hostname is required for SFTP")

        logger.debug(f"Uploading to SFTP: {server_config.hostname}:{ssh_port} {remote_path}")

        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.ssh.connect(
            hostname=hostname,
            port=ssh_port,
            username=server_config.username,
            password=connect_password,
            pkey=private_key,
            look_for_keys=False,
            allow_agent=False,
        )
        with self.ssh.open_sftp() as sftp:
            sftp.putfo(data_to_upload, remote_path)
        self.ssh.close()
